package com.shujia.spark.mllib

import org.apache.spark.ml.classification.{NaiveBayes, NaiveBayesModel}
import org.apache.spark.ml.feature.{HashingTF, IDF, IDFModel, Tokenizer}
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo8TextClass {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[8]")
      .appName("text")
      .getOrCreate()


    import spark.implicits._
    import org.apache.spark.sql.functions._

    //读取数据
    val testDF: DataFrame = spark
      .read
      .format("csv")
      .option("sep", "\t")
      .schema("label DOUBLE, text STRING")
      .load("data/text.txt")
      .repartition(8) //修改DF的分区


    //对中文文本进行分词
    val segment: UserDefinedFunction = udf((text: String) => {
      //中文分词，按空格拼接，返回
      IKUtil.fit(text).mkString(" ")
    })


    /**
     * 中文分词，过滤脏数据
     */
    val sentenceDF: Dataset[Row] = testDF
      //中文分词
      .select($"label", segment($"text") as "sentence")
      //过滤脏数据
      .where(size(split($"sentence", " ")) > 1)

    /**
     * 使用英文分词器预处理数据，按照空格分割
     */
    val tokenizer: Tokenizer = new Tokenizer()
      .setInputCol("sentence") //输入列
      .setOutputCol("words") //输出列

    val wordsDF: DataFrame = tokenizer.transform(sentenceDF)


    /**
     * 增加词频 -- TF
     */

    val hashingTF: HashingTF = new HashingTF()
      .setInputCol("words") //输入列
      .setOutputCol("rawFeatures") //输出列
    //.setNumFeatures(20) //总的词语的数量

    val tfDF: DataFrame = hashingTF.transform(wordsDF)

    tfDF.cache()

    /**
     * 增加逆文本频率-- IDF
     */

    val idf: IDF = new IDF()
      .setInputCol("rawFeatures") //输入列
      .setOutputCol("features") //输出列

    //训练模型
    val iDFModel: IDFModel = idf.fit(tfDF)

    val idfDF: DataFrame = iDFModel.transform(tfDF)

    /**
     * 将数据拆分成训练集和测试集
     */
    val Array(train, test) = idfDF.randomSplit(Array(0.8, 0.2))

    /**
     * 选择算法
     */
    val naiveBayes = new NaiveBayes()

    //将训练集带入算法训练模型
    val model: NaiveBayesModel = naiveBayes.fit(train)

    //使用测试集测试模型的准确率
    val frame: DataFrame = model.transform(test)

    //准确率
    val acc: Double = frame.where($"label" === $"prediction").count().toDouble /
      frame.count()

    println(s"准确率：$acc")

    //精确率
    val p: Double = frame.where($"prediction" === 1.0 and $"label" === $"prediction").count().toDouble /
      frame.where($"prediction" === 1.0).count()

    println(s"精确率：$p")

    //召回率
    val c: Double = frame.where($"prediction" === 1.0 and $"label" === $"prediction").count().toDouble /
      frame.where($"label" === 1.0).count()

    println(s"召回率:$c")


  }

}
